import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import yfinance as yf
import seaborn as sns
from sklearn import linear_model
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.ensemble import RandomForestClassifier
from sklearn.neural_network import MLPClassifier
import datetime
import backtrader as bt
from backtrader.feeds import PandasData
import backtrader.analyzers as btanalyzers

plt.style.use('seaborn-colorblind')

ticker = 'TSLA'
start = datetime.datetime(2010, 1, 1)
end = datetime.datetime(2020, 12, 30)
stock = yf.download(ticker, progress=True, actions=True, start=start, end=end)
stock = stock['Adj Close']
stock = pd.DataFrame(stock)
stock.rename(columns={"Adj Close": ticker}, inplace=True)
stock['returns'] = np.log(stock/stock.shift(1))
stock.dropna(inplace=True)
stock['direction'] = np.sign(stock['returns']).astype(int)

lags = [1, 2,3, 4,5]
cols = []
for lag in lags:
    col = f'rnt_lag{lag}'
    stock[col] = stock['returns'].shift(lag)
    cols.append(col)
stock.dropna(inplace=True)


def create_bins(data, bins=[0]):
    global cols_bin
    cols_bin = []
    for col in cols:
        col_bin = col + '_bin'
        data[col_bin] = np.digitize(data[col], bins=bins)
        cols_bin.append(col_bin)

create_bins(stock)

models = {
    'log_reg': linear_model.LogisticRegression(),
    'gauss_nb': GaussianNB(),
    'svm': SVC(),
    'random_forest': RandomForestClassifier(max_depth=10, n_estimators=100),
    'MLP': MLPClassifier(max_iter=100)
}

def fit_models(data):
    mfit = {model: models[model].fit(data[cols_bin], data['direction']) for model in models.keys()}

def derive_positions(data):
    for model in models.keys():
        data['pos_' + model] = models[model].predict(data[cols_bin])

def evaluate_strats(data):
    global strategy_rtn
    strategy_rtn = []
    for model in models.keys():
        col = 'strategy_' + model
        data[col] = data['pos_'+ model] * data['returns']
        strategy_rtn.append(col)
    strategy_rtn.insert(0, 'returns')

fit_models(stock)
derive_positions(stock)
evaluate_strats(stock)

print(stock[strategy_rtn].sum().apply(np.exp))



price = yf.download(ticker, progress=True, actions=True, start=start,end=end)
predictions = stock['strategy_random_forest']
predictions = pd.DataFrame(predictions)
predictions.rename(columns={'strategy_random_forest': 'predicted'}, inplace=True)
price = predictions.join(price, how='right').dropna()
OHLCH=['open', 'high', 'low', 'close', 'volume']

class SignalData(PandasData):
    cols = OHLCH + ['predicted']
    lines = tuple(cols)
    params = {c: -1 for c in cols}
    params.update({'datetime': None})
    params = tuple(params.items())

class RandomForestStrategy(bt.Strategy):
    params = dict(
    
    )
    
    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.data.datetime[0]
        if isinstance(dt, float):
            dt = bt.num2date(dt)
        print('%s, %s' % (dt.isoformat(), txt))
    
    def __init__(self):
        self.data_predicted = self.datas[0].predicted
        self.data_open = self.datas[0].open
        self.data_close = self.datas[0].close
        
        self.order = None
        self.price = None
        self.comm = None
        
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log('BUY EXECUTED, %.2f' % order.executed.price)
            elif order.issell():
                self.log('SELL EXECUTED, %.2f' % order.executed.price)
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')
        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
                 (trade.pnl, trade.pnlcomm))
    
    def next_open(self):
        if not self.position:
            if self.data_predicted > 0:
                size = int((self.broker.getcash()/self.datas[0].open)*0.9)
                self.log(f"BUY CREATED - SIZE : {size}, CASH: {self.broker.getcash():.2f}, OPEN: {self.data_open[0]}, CLOSE: {self.data_close[0]}")
                self.buy(size=size)
        else:
            if self.data_predicted < 0:
                self.log(f'SELL CREATED - SIZE: {self.position.size}')
                self.close()

data = SignalData(dataname=price)
cerebro = bt.Cerebro(cheat_on_open=True)
cerebro.addstrategy(RandomForestStrategy)
cerebro.adddata(data, name=ticker)
cerebro.broker.setcash(10000.0)
cerebro.broker.setcommission(commission=0.0003)

cerebro.addanalyzer(btanalyzers.SharpeRatio, _name="sharpe")
cerebro.addanalyzer(btanalyzers.DrawDown, _name="drawdown")
cerebro.addanalyzer(btanalyzers.Returns, _name="returns")

print('the beigin value is {:.2f}'.format(cerebro.broker.getvalue()))

back = cerebro.run()

print('the end value is {:.2f}'.format(cerebro.broker.getvalue()))


par_list = [[x.analyzers.returns.get_analysis()['rtot'],
             x.analyzers.returns.get_analysis()['rnorm100'],
             x.analyzers.drawdown.get_analysis()['max']['drawdown'],
             x.analyzers.sharpe.get_analysis()['sharperatio']
            ]for x in back
]
par_df = pd.DataFrame(par_list, columns=['Total returns', 'APR', 'drawdown', 'sharperatio'])
print(par_df)
cerebro.plot(style='candle')